/**
 * @file random_util.c
 *
 * @date Sep 22, 2011
 * @author seth
 *
 * @brief This file is part of EBEX State Estimator, created for the EBEX project
 *
 * This software is copyright (C) 2011 Columbia University
 *
 * EBEX State Estimator is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * EBEX State Estimator is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with EBEX State Estimator; if not, write to the Free Software Foundation, Inc.,
 * 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 */

#include <mkl.h>
#include <stdbool.h>

#include <log_util.h>
#include <random_util.h>

static volatile VSLStreamStatePtr current_state = NULL;
static volatile random_generator_t generator = RANDOM_MT19937;

random_generator_t ran_get_generator(void)
{
	return generator;
}

bool ran_set_generator(random_generator_t m_generator)
{
	if (generator != m_generator)
	{
		generator = m_generator;
		if (current_state) ran_init();
	}
	return true;
}

bool ran_seed(const unsigned int seed)
{
	int retval;
	int delay = 7;
	VSLStreamStatePtr temp_ptr = NULL;

	retval = vslNewStream(&temp_ptr, generator, seed);
	switch (retval)
	{
		case VSL_RNG_ERROR_INVALID_BRNG_INDEX:
			elog(log_level_fatal, "Invalid generator %u", (unsigned) generator);
			return false;
		case VSL_ERROR_MEM_FAILURE:
			die("Out of memory!");
			return false;
	}

	EBEX_SWAP(temp_ptr, current_state);

	if (temp_ptr)
		vslDeleteStream(&temp_ptr);

	/// Warm it up a bit
	while (delay--)
	{
		(void)ran_int_uniform();
	}
	return true;
}

void ran_init(void)
{
	int delay = 7;
	ran_seed((unsigned)time(NULL));

	/// Warm it up a bit
	while (delay--)
	{
		(void)ran_int_uniform();
	}
}

void ran_gamma_n(const double m_shape, const double m_offset, const double m_scale, double *m_output, const int m_size)
{
	int retval;
	retval = vdRngGamma(VSL_RNG_METHOD_GAMMA_GNORM, current_state, m_size, m_output, m_shape, m_offset, m_scale);
	switch (retval)
	{
		case VSL_STATUS_OK:
			break;
		case VSL_ERROR_NULL_PTR:
			elog(log_level_fatal, "Passed NULL pointer");
			return;
		case VSL_RNG_ERROR_BAD_STREAM:
		case VSL_RNG_ERROR_BAD_UPDATE:
		case VSL_RNG_ERROR_NO_NUMBERS:
			elog(log_level_fatal, "Error in RNG stream");
			return;
		case VSL_RNG_ERROR_QRNG_PERIOD_ELAPSED:
			elog(log_level_fatal, "Eclipsed RNG period");
			break;
	}
}

/**
 * Returns an integer distributed uniformly over [min,max)
 * @param min Minimum, inclusive value for distribution
 * @param max Maximum, non-inclusive value for distribution
 * @return Uniform, random integer
 */
int ran_int_uniform_range(const int min, const int max)
{
	int retval = max;

	if (max - min)
		if (viRngUniform(VSL_RNG_METHOD_UNIFORM_STD, current_state, 1, &retval, min, max) != VSL_STATUS_OK) die("Error in RNG");
	return retval;
}

void ran_double_range_uniform_n(const double m_min, const double m_max, double *m_output, size_t m_num)
{
	if (m_max > m_min)
	{
		if (vdRngUniform(VSL_RNG_METHOD_UNIFORM_STD, current_state, m_num, m_output, m_min, m_max) != VSL_STATUS_OK) die("Error in RNG");
	}
	else
	{
		for (size_t i = 0; i < m_num; i++) m_output[i] = m_max;
	}
}

void ran_gaussian_n(const double m_mean, const double m_stddev, double *m_output, int m_num)
{
	if (vdRngGaussian(VSL_RNG_METHOD_GAUSSIAN_ICDF, current_state, m_num,
			m_output, m_mean, m_stddev) != VSL_STATUS_OK) die("Error in RNG");
}

void ran_gaussian_md(const double *m_mean, const double *m_cov, double *m_output, int m_dim, int m_num)
{
	if (vdRngGaussianMV(VSL_RNG_METHOD_GAUSSIANMV_ICDF, current_state, m_num,
			m_output, m_dim, VSL_MATRIX_STORAGE_FULL, m_mean, m_cov) != VSL_STATUS_OK) die("Error in RNG");
}

void ran_binomial_n(const int m_trials, const double m_probability, int *m_output, const int m_num)
{
	if (viRngBinomial(VSL_RNG_METHOD_BINOMIAL_BTPE, current_state, m_num, m_output, m_trials, m_probability) != VSL_STATUS_OK) die("Error in RNG");
}

void ran_weibull_n(const double m_shape, const double m_displacement, const double m_scale, double *m_output, int m_num)
{
	if (vdRngWeibull(VSL_RNG_METHOD_WEIBULL_ICDF, current_state, m_num, m_output, m_shape, m_displacement, m_scale) != VSL_STATUS_OK) die("Error in RNG");
}


void ran_int_permutation(const int m_size, int *m_inputarray, int *m_outputarray)
{
	int i, j = 0; /* Loop variables over arrays. */
	int pos; /* Randomly selected index. */

	if (!m_inputarray || !m_outputarray)
		die("NULL pointer to int array passed.");

	for (i = m_size - 1; i > 0; i--)
	{
		pos = ran_uint_uniform(i);
		m_outputarray[j++] = m_inputarray[pos];
		m_inputarray[pos] = m_inputarray[i];
	}

	m_outputarray[j] = m_inputarray[0];

	return;
}

double ran_sample_mean(random_sample_t *m_param)
{
	double retval = nan("INVALID");

	switch(m_param->type)
	{
		case random_dist_beta:
			retval = (m_param->beta.a) / (m_param->beta.a + m_param->beta.b);
			break;
		case random_dist_binomial:
			retval = m_param->binomial.trials * m_param->binomial.probability;
			break;
		case random_dist_dirichlet:
			retval = 1.0;
			break;
		case random_dist_epanechnikov:
			retval = 0.5;
			break;
		case random_dist_gamma:
			retval = m_param->gamma.shape * m_param->gamma.scale;
			break;
		case random_dist_gaussian:
			retval = m_param->gaussian.mean;
			break;
		case random_dist_pearsonii:
			retval = 0.0;
			break;
		case random_dist_uniform:
			retval = 0.0;
			break;
		case random_dist_weibull:
			retval = m_param->weibull.scale * tgamma(1.0 + 1.0 / m_param->weibull.shape);
			break;
		default:
			die("Invalid distribution type");
			break;
	}
	return retval * m_param->scale + m_param->offset;
}

double ran_sample_mode(random_sample_t *m_param)
{
	double retval = nan("INVALID");

	switch(m_param->type)
	{
		case random_dist_beta:
			if (m_param->beta.a > 0.0 && m_param->beta.b > 0.0)
				retval = (m_param->beta.a - 1.0) / (m_param->beta.a + m_param->beta.b - 2.0);
			break;
		case random_dist_binomial:
			retval = (m_param->binomial.trials + 1.0) * m_param->binomial.probability;
			break;
		case random_dist_dirichlet:
			retval = 1.0;
			break;
		case random_dist_epanechnikov:
			retval = 0.5;
			break;
		case random_dist_gamma:
			if (m_param->gamma.shape >= 1.0)
				retval = (m_param->gamma.shape - 1.0) * m_param->gamma.scale;
			break;
		case random_dist_gaussian:
			retval = m_param->gaussian.mean;
			break;
		case random_dist_pearsonii:
			retval = 0.0;
			break;
		case random_dist_uniform:
			retval = 0.0;
			break;
		case random_dist_weibull:
			if (m_param->weibull.shape > 1.0)
				retval = m_param->weibull.scale * pow((m_param->weibull.shape - 1.0)/m_param->weibull.shape, 1.0 / m_param->weibull.shape);
			break;
		default:
			die("Invalid distribution type");
			break;
	}
	return retval * m_param->scale + m_param->offset;
}
void ran_sample_scale_offset_n(double *m_dest, const random_sample_t *m_param, double m_scale, double m_offset, size_t m_num, size_t m_stride)
{
	int binomial_out;
	double *temp_dest = m_dest;

	if (m_stride != 1)
	{
		temp_dest = e_malloc(m_num * sizeof(double));
	}
	switch(m_param->type)
	{
		case random_dist_none:
			e_memset(temp_dest, 0, m_num * sizeof(double));
			break;
		case random_dist_beta:
			ran_beta_n(m_param->beta.a, m_param->beta.b, temp_dest, m_num);
			break;
		case random_dist_binomial:
			for (size_t i = 0; i < m_num; i++)
			{
				ran_binomial_n((int)m_param->binomial.trials, m_param->binomial.probability, &binomial_out, 1);
				temp_dest[i] = (double)binomial_out / (double)m_param->binomial.probability;
			}
			break;
		case random_dist_dirichlet:
			ran_dirichlet_n(m_param->dirichlet.shape, temp_dest, 1, m_num);
			break;
		case random_dist_epanechnikov:
			ran_epanechnikov_n(temp_dest, m_num);
			break;
		case random_dist_gamma:
			ran_gamma_n(m_param->gamma.shape, 0.0, m_param->gamma.scale, temp_dest, m_num);
			break;
		case random_dist_gaussian:
			ran_gaussian_n(m_param->gaussian.mean, m_param->gaussian.stddev, temp_dest, m_num);
			break;
		case random_dist_pearsonii:
			for (size_t i = 0; i < m_num; i++)
			{
				ran_mvpearson_ii_md(&temp_dest[i], (const double[]){0.0}, 1);
			}
			break;
		case random_dist_uniform:
			ran_double_range_uniform_n(m_param->uniform.min, m_param->uniform.max, temp_dest, m_num);
			break;
		case random_dist_weibull:
			ran_weibull_n(m_param->weibull.shape, 0.0, m_param->weibull.scale, temp_dest, m_num);
			break;
		default:
			die("Invalid distribution type");
			break;
	}
	if (!ISDBL_TINY(m_scale - 1.0))
	{
		for (size_t i = 0; i < m_num; i++) temp_dest[i] *= m_scale;
	}
	if (!ISDBL_TINY(m_offset))
	{
		for (size_t i = 0; i < m_num; i++) temp_dest[i] += m_offset;
	}

	if (m_stride != 1)
	{
		for (size_t i = 0; i < m_num; i++)
		{
			m_dest[i * m_stride] = temp_dest[i];
		}
		e_free(temp_dest);
	}
}

double ran_sample_scale_offset(const random_sample_t *m_param, double m_scale, double m_offset)
{
	double retval;
	int binomial_out;

	switch(m_param->type)
	{
		case random_dist_none:
			retval = 0.0;
			break;
		case random_dist_beta:
			retval = ran_beta(m_param->beta.a, m_param->beta.b);
			break;
		case random_dist_binomial:
			ran_binomial_n((int)m_param->binomial.trials, m_param->binomial.probability, &binomial_out, 1);
			retval = (double)binomial_out / (double)m_param->binomial.probability;
			break;
		case random_dist_dirichlet:
			ran_dirichlet(m_param->dirichlet.shape, (double[]) {0.0}, &retval, 1);
			break;
		case random_dist_epanechnikov:
			retval = ran_epanechnikov();
			break;
		case random_dist_gamma:
			retval = ran_gamma(m_param->gamma.shape, 0.0, m_param->gamma.scale);
			break;
		case random_dist_gaussian:
			retval = ran_gaussian(m_param->gaussian.mean, m_param->gaussian.stddev);
			break;
		case random_dist_pearsonii:
			ran_mvpearson_ii_md(&retval, (const double[]){0.0}, 1);
			break;
		case random_dist_uniform:
			retval = ran_double_range_uniform(m_param->uniform.min, m_param->uniform.max);
			break;
		case random_dist_weibull:
			ran_weibull_n(m_param->weibull.shape, 0.0, m_param->weibull.scale, &retval, 1);
			break;
		default:
			die("Invalid distribution type");
			break;
	}
	return retval * m_scale + m_offset;
}

void ran_sample_scale_offset_3d(const random_sample_3d_t *m_param, double *m_retvals, const double *m_scale, const double *m_offset)
{
	double cov[3][3] =
			{
			{1.0, 0.0, 0.0},
			{0.0, 1.0, 0.0},
			{0.0, 0.0, 1.0}
			};

	switch(m_param->type)
	{
		case random_dist_none:
			m_retvals[0] = 0.0;
			m_retvals[1] = 0.0;
			m_retvals[2] = 0.0;
			break;
		case random_dist_dirichlet:
			ran_dirichlet(m_param->dirichlet[0].shape, m_offset, m_retvals, 3);
			break;
		case random_dist_gaussian:
			cov[0][0] = pow(m_param->gaussian[0].stddev, 2.0);
			cov[1][1] = pow(m_param->gaussian[1].stddev, 2.0);
			cov[2][2] = pow(m_param->gaussian[2].stddev, 2.0);
			ran_gaussian_md(m_offset, cov[0], m_retvals, 3, 1);
			break;
		case random_dist_pearsonii:
			ran_mvpearson_ii_md(m_retvals, m_offset, 3);
			break;
		case random_dist_uniform:
		case random_dist_spherical:
			ran_unit_sphere_md(m_retvals, m_scale[0],(double[]){0.0,0.0,0.0}, 3);
			break;
		default:
			die("Invalid distribution type");
			break;
	}
}
